use std::sync::Arc;

use async_trait::async_trait;
use bytestring::ByteString;
use datafusion::arrow::array::RecordBatch;
use datafusion::prelude::{DataFrame, SessionContext};
use datafusion::sql::sqlparser::ast::{Ident, ObjectName, ObjectType, ShowStatementFilter};
use either::Either;

use common::{Error, LResult};
use engine::{DB, Storage};
use function::register_functions;
use mem_store::MemEngineFactory;

pub struct LiteDB {
    session_ctx: Arc<SessionContext>,
    storage: Storage,
}

impl LiteDB {
    #[allow(dead_code)]
    pub fn create() -> LiteDB {
        let ctx = Arc::new(SessionContext::new());
        register_functions(ctx.clone());
        let mut storage = Storage::create(ctx.clone());
        let mem_engine = MemEngineFactory::create();
        storage.register_engine_factory(mem_engine);
        Self {
            session_ctx: ctx,
            storage,
        }
    }

    #[allow(dead_code)]
    pub async fn execute_query(&self, sql: &str) -> LResult<Vec<RecordBatch>> {
        let data = self.query(sql).await?;
        match data {
            Either::Left(df) => {
                df.collect().await.map_err(|e| Error::ExecutionErr(format!("query {} failed: {:?}", sql, e)))
            }
            Either::Right(r) => {
                Ok(vec![r])
            }
        }
    }
}

#[async_trait]
impl DB for LiteDB {
    #[inline]
    async fn insert(&self, line: ByteString) -> LResult<()> {
        self.storage.insert(line).await
    }

    async fn show_tables(&self, extended: bool, full: bool, db_name: &Option<Ident>, filter: &Option<ShowStatementFilter>) -> LResult<RecordBatch> {
        self.storage.show_tables(extended, full, db_name, filter).await
    }

    async fn drop(&self, object_type: &ObjectType, if_exists: bool, name: &Vec<ObjectName>) -> LResult<RecordBatch> {
        self.storage.drop(object_type, if_exists, name).await
    }

    async fn create_database(&self, db_name: &ObjectName, if_not_exists: bool) -> LResult<RecordBatch> {
        self.storage.create_database(db_name, if_not_exists).await
    }

    fn session_context(&self) -> &SessionContext {
        &*self.session_ctx
    }

    fn post_query(&self, df: DataFrame) -> LResult<DataFrame> {
        self.storage.post_transform_query(df)
    }
}


#[cfg(test)]
mod test {
    use datafusion::arrow::array::{Array, GenericStringArray, RecordBatch, UInt64Array};

    use common::LResult;
    use engine::DB;

    use crate::db::LiteDB;

    #[tokio::test]
    async fn test() -> LResult<()> {
        let mut db = LiteDB::create();
        insert_data(&mut db).await?;
        assert_sql_one_batch(&db, "select count(ts_s) from test", |b| {
            assert_eq!(1, b.columns().len());
            assert_eq!(1, b.column(0).len());
            let a = unsafe { b.column(0).as_any().downcast_ref_unchecked::<UInt64Array>() };
            assert_eq!(2, a.value(0));
        }).await;
        assert_sql_one_batch(&db, "select $t_tag_a from test", |b| {
            assert_eq!(1, b.columns().len());
            assert_eq!(2, b.column(0).len());
            let tag_column = b.column(0).as_any().downcast_ref::<GenericStringArray<i32>>().unwrap();
            assert_eq!("abc", tag_column.value(0));
            assert_eq!("abc", tag_column.value(1));
        }).await;
        Ok(())
    }

    async fn insert_data(db: &mut LiteDB) -> LResult<()> {
        {
            let e0 = "test,tag_a=abc,tag_b=aaa field_a=123,field_b=\"aaa\"";
            db.insert(e0.into()).await?;
        }
        {
            let e0 = "test,tag_a=abc,tag_b=bbb field_a=222,field_b=\"bbb\"";
            db.insert(e0.into()).await?;
        }
        Ok(())
    }

    async fn assert_sql_one_batch<F: FnOnce(&RecordBatch)>(db: &LiteDB, sql: &str, f: F) {
        let v = db.execute_query(sql).await.unwrap();
        assert_eq!(1, v.len(), "assert_sql_one_batch has {} recordBatch", v.len());
        f(&v[0])
    }
}